What does this PR do?

Moshi is the latest Kyutai model. It is a streaming speech-to-speech model, that can also do an inner dialogue (i.e it outputs text as well).

In particular, it means that Moshi deals with 3 streams of information:

1. The user's audio
2. Moshi's audio
3. Moshi's textual output

Similarly to Musicgen, audio is represented with audio codebooks, which can be interpreted like tokens. The main difference between text tokens and audio codebooks is that audio codebooks introduce an additional dimension of information.
Text tokens are typically of dim (batch_size, sequence_length) but audio tokens are of dim (batch_size, num_codebooks, sequence_length).

It's made of 3 components:

1. The main decoder (Helium in the paper)

Here, it corresponds to MoshiForCausalLM. It is strictly a classic text LLM, that uses an architecture similar to Gemma. In other words, it takes text tokens, embeds them, pass them through the decoder and a language head, to get text logits.

2. The depth decoder

On its own, it's also a classic LLM, but this time, instead of generating over the time dimension, it generates over the codebook dimension.

It also means that its context length is num_codebooks -> it can't generate more than num_codebooks.

Another interesting difference with a classic LLM is that each timestamp (here it correspond to each codebook) got its own set of Linear Layers and Embeddings.

3. Mimi

It's the audio encoder from Kyutai, that has recently been integrated to transformers, which is used to "tokenize" audio. It has the same use that Encodec has in Musicgen.

Architecture choice:

1. MoshiForCausalLM corresponds to the main decoder, it can be used as a textual LLM.

2. MoshiDepthDecoder is the depth decoder mentioned above

3. MoshiForConditionalGeneration encapsulates the main decoder, the depth decoder and the audio encoder.

Conceptually, MoshiForConditionalGeneration takes as input one stream of text and two streams of audio inputs - what the user has said so far, and what the model have generated so far - and generates two streams - a text stream and an audio stream.

How does it work:

-> The input streams are embedded and combined into inputs_embeds.

-> inputs_embeds is passed through the main decoder. There's nothing special done here, it's the same operation as Gemma or so on.

-> The main decoder outputs text logits but also its last hidden state which is called temporal context in the picture above.

-> the depth decoder switches the dimension on which we generate (codebooks instead of time). It uses the token generated from text logits and the temporal context to auto-regressively generate audio codebooks.